{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom LLM\n",
    "This notebook goes over how to create a custom LLM wrapper, in case you want to use your own LLM or a different wrapper than one that is supported in LangChain.<br>\n",
    "本笔记本介绍了如何创建自定义LLM包装器，以防您想要使用自己的LLM包装器或与LangChain中支持的包装器不同的包装器。\n",
    "\n",
    "Wrapping your LLM with the standard LLM interface allow you to use your LLM in existing LangChain programs with minimal code modifications!<br>\n",
    "用标准 LLM 接口包装你LLM，允许你在现有的LangChain程序中使用你的LLM，只需最少的代码修改！\n",
    "\n",
    "As an bonus, your LLM will automatically become a LangChain Runnable and will benefit from some optimizations out of the box, async support, the astream_events API, etc.<br>\n",
    "作为奖励，您将LLM自动成为LangChain Runnable ，并将受益于一些开箱即用的优化，异步支持， astream_events API等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict, Iterator, List, Mapping, Optional\n",
    "\n",
    "from langchain_core.callbacks.manager import CallbackManagerForLLMRun\n",
    "from langchain_core.language_models.llms import LLM\n",
    "from langchain_core.outputs import GenerationChunk\n",
    "\n",
    "\n",
    "class CustomLLM(LLM):\n",
    "    \"\"\"A custom chat model that echoes the first `n` characters of the input.\n",
    "\n",
    "    When contributing an implementation to LangChain, carefully document\n",
    "    the model including the initialization parameters, include\n",
    "    an example of how to initialize the model and include any relevant\n",
    "    links to the underlying models documentation or API.\n",
    "\n",
    "    Example:\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            model = CustomChatModel(n=2)\n",
    "            result = model.invoke([HumanMessage(content=\"hello\")])\n",
    "            result = model.batch([[HumanMessage(content=\"hello\")],\n",
    "                                 [HumanMessage(content=\"world\")]])\n",
    "    \"\"\"\n",
    "\n",
    "    n: int\n",
    "    \"\"\"The number of characters from the last message of the prompt to be echoed.\"\"\"\n",
    "\n",
    "    def _call(\n",
    "        self,\n",
    "        prompt: str,\n",
    "        stop: Optional[List[str]] = None,\n",
    "        run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> str:\n",
    "        \"\"\"Run the LLM on the given input.\n",
    "\n",
    "        Override this method to implement the LLM logic.\n",
    "\n",
    "        Args:\n",
    "            prompt: The prompt to generate from.\n",
    "            stop: Stop words to use when generating. Model output is cut off at the\n",
    "                first occurrence of any of the stop substrings.\n",
    "                If stop tokens are not supported consider raising NotImplementedError.\n",
    "            run_manager: Callback manager for the run.\n",
    "            **kwargs: Arbitrary additional keyword arguments. These are usually passed\n",
    "                to the model provider API call.\n",
    "\n",
    "        Returns:\n",
    "            The model output as a string. Actual completions SHOULD NOT include the prompt.\n",
    "        \"\"\"\n",
    "        if stop is not None:\n",
    "            raise ValueError(\"stop kwargs are not permitted.\")\n",
    "        return probmpt[: self.n]\n",
    "\n",
    "    def _stream(\n",
    "        self,\n",
    "        prompt: str,\n",
    "        stop: Optional[List[str]] = None,\n",
    "        run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Iterator[GenerationChunk]:\n",
    "        \"\"\"Stream the LLM on the given prompt.\n",
    "\n",
    "        This method should be overridden by subclasses that support streaming.\n",
    "\n",
    "        If not implemented, the default behavior of calls to stream will be to\n",
    "        fallback to the non-streaming version of the model and return\n",
    "        the output as a single chunk.\n",
    "\n",
    "        Args:\n",
    "            prompt: The prompt to generate from.\n",
    "            stop: Stop words to use when generating. Model output is cut off at the\n",
    "                first occurrence of any of these substrings.\n",
    "            run_manager: Callback manager for the run.\n",
    "            **kwargs: Arbitrary additional keyword arguments. These are usually passed\n",
    "                to the model provider API call.\n",
    "\n",
    "        Returns:\n",
    "            An iterator of GenerationChunks.\n",
    "        \"\"\"\n",
    "        for char in prompt[: self.n]:\n",
    "            chunk = GenerationChunk(text=char)\n",
    "            if run_manager:\n",
    "                run_manager.on_llm_new_token(chunk.text, chunk=chunk)\n",
    "\n",
    "            yield chunk\n",
    "\n",
    "    @property\n",
    "    def _identifying_params(self) -> Dict[str, Any]:\n",
    "        \"\"\"Return a dictionary of identifying parameters.\"\"\"\n",
    "        return {\n",
    "            # The model name allows users to specify custom token counting\n",
    "            # rules in LLM monitoring applications (e.g., in LangSmith users\n",
    "            # can provide per token pricing for their model and monitor\n",
    "            # costs for the given LLM.)\n",
    "            \"model_name\": \"CustomChatModel\",\n",
    "        }\n",
    "\n",
    "    @property\n",
    "    def _llm_type(self) -> str:\n",
    "        \"\"\"Get the type of language model used by this chat model. Used for logging purposes only.\"\"\"\n",
    "        return \"custom\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mCustomLLM\u001b[0m\n",
      "Params: {'model_name': 'CustomChatModel'}\n"
     ]
    }
   ],
   "source": [
    "llm = CustomLLM(n=5)\n",
    "print(llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This '"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm.invoke(\"This is a foobar thing\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"you are a bot\"), (\"human\", \"{input}\")]\n",
    ")\n",
    "llm = CustomLLM(n=7)\n",
    "chain = prompt | llmb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_chain_start', 'run_id': '7de3c8da-3c6f-49df-ab70-ad4b98f35197', 'name': 'RunnableSequence', 'tags': [], 'metadata': {}, 'data': {'input': {'input': 'hello there!'}}}\n",
      "{'event': 'on_prompt_start', 'name': 'ChatPromptTemplate', 'run_id': '2bc253f7-6a35-4d25-88fd-275f69234259', 'tags': ['seq:step:1'], 'metadata': {}, 'data': {'input': {'input': 'hello there!'}}}\n",
      "{'event': 'on_prompt_end', 'name': 'ChatPromptTemplate', 'run_id': '2bc253f7-6a35-4d25-88fd-275f69234259', 'tags': ['seq:step:1'], 'metadata': {}, 'data': {'input': {'input': 'hello there!'}, 'output': ChatPromptValue(messages=[SystemMessage(content='you are a bot'), HumanMessage(content='hello there!')])}}\n",
      "{'event': 'on_llm_start', 'name': 'CustomLLM', 'run_id': '2304042d-1047-4432-adb0-05c97c657fcd', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'input': {'prompts': ['System: you are a bot\\nHuman: hello there!']}}}\n",
      "{'event': 'on_llm_stream', 'name': 'CustomLLM', 'run_id': '2304042d-1047-4432-adb0-05c97c657fcd', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': 'S'}}\n",
      "{'event': 'on_chain_stream', 'run_id': '7de3c8da-3c6f-49df-ab70-ad4b98f35197', 'tags': [], 'metadata': {}, 'name': 'RunnableSequence', 'data': {'chunk': 'S'}}\n",
      "{'event': 'on_llm_stream', 'name': 'CustomLLM', 'run_id': '2304042d-1047-4432-adb0-05c97c657fcd', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': 'y'}}\n",
      "{'event': 'on_chain_stream', 'run_id': '7de3c8da-3c6f-49df-ab70-ad4b98f35197', 'tags': [], 'metadata': {}, 'name': 'RunnableSequence', 'data': {'chunk': 'y'}}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\Dev\\anaconda3\\envs\\langchain0_1\\Lib\\site-packages\\langchain_core\\_api\\beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.\n",
      "  warn_beta(\n"
     ]
    }
   ],
   "source": [
    "idx = 0\n",
    "async for event in chain.astream_events({\"input\": \"hello there!\"}, version=\"v1\"):\n",
    "    print(event)\n",
    "    idx += 1\n",
    "    if idx > 7:\n",
    "        # Truncate\n",
    "        break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
